//! Implementation of chain initialization for the Shell
use std::collections::BTreeMap;
use std::ops::ControlFlow;

use masp_primitives::merkle_tree::CommitmentTree;
use masp_primitives::sapling::Node;
use masp_proofs::bls12_381;
use namada_sdk::account::protocol_pk_key;
use namada_sdk::collections::HashMap;
use namada_sdk::eth_bridge::EthBridgeStatus;
use namada_sdk::hash::Hash as CodeHash;
use namada_sdk::parameters::Parameters;
use namada_sdk::proof_of_stake::{self, BecomeValidator, PosParams};
use namada_sdk::state::StorageWrite;
use namada_sdk::time::{TimeZone, Utc};
use namada_sdk::token::storage_key::masp_token_map_key;
use namada_sdk::token::{credit_tokens, write_denom};
use namada_sdk::{eth_bridge, ibc};
use namada_vm::validate_untrusted_wasm;

use super::*;
use crate::config::genesis::chain::{
    FinalizedEstablishedAccountTx, FinalizedTokenConfig,
    FinalizedValidatorAccountTx,
};
use crate::config::genesis::templates::{TokenBalances, TokenConfig};
use crate::config::genesis::transactions::{
    BondTx, EstablishedAccountTx, Signed as SignedTx, ValidatorAccountTx,
};
use crate::tendermint_proto::google::protobuf;
use crate::wasm_loader;

/// Errors that represent panics in normal flow but get demoted to errors
/// when dry-running genesis files in order to accumulate as many problems
/// as possible in a report.
#[derive(Error, Debug, Clone, PartialEq)]
enum Panic {
    #[error(
        "No VP found matching the expected implicit VP sha256 hash: \
         {0}\n(this will be `None` if no wasm file was found for the implicit \
         vp)"
    )]
    MissingImplicitVP(String),
    #[error("Missing validity predicate for {0}")]
    MissingVpWasmConfig(String),
    #[error("Could not find checksums.json file")]
    ChecksumsFile,
    #[error("Invalid wasm code sha256 hash for {0}")]
    Checksum(String),
    #[error(
        "Config for token '{0}' with configured balance not found in genesis"
    )]
    MissingTokenConfig(String),
    #[error("The MASP parameters for the native token is missing")]
    MissingMaspParams,
    #[error("Failed to read wasm {0} with reason: {1}")]
    ReadingWasm(String, String),
}

/// Warnings generated by problems in genesis files.
#[derive(Error, Debug, PartialEq)]
enum Warning {
    #[error("The wasm {0} isn't allowed.")]
    DisallowedWasm(String),
    #[error("Genesis init genesis validator tx for {0} failed with {1}.")]
    Validator(String, String),
    #[error(
        "Genesis bond by {0} to validiator {1} of {2} NAM failed with reason: \
         {3}"
    )]
    FailedBond(String, String, token::DenominatedAmount, String),
}

impl<D, H> Shell<D, H>
where
    D: DB + for<'iter> DBIter<'iter> + Sync + 'static,
    H: StorageHasher + Sync + 'static,
{
    /// Create a new genesis for the chain with specified id. This includes
    /// 1. A set of initial users and tokens
    /// 2. Setting up the validity predicates for both users and tokens
    /// 3. Validators
    /// 4. The PoS system
    /// 5. The Ethereum bridge parameters
    ///
    /// INVARIANT: This method must not commit the state changes to DB.
    pub fn init_chain(
        &mut self,
        init: request::InitChain,
        #[cfg(any(test, feature = "testing", feature = "benches"))]
        num_validators: u64,
    ) -> ShellResult<response::InitChain> {
        let mut response = response::InitChain::default();
        let chain_id = self.state.in_mem().chain_id.as_str();
        if chain_id != init.chain_id.as_str() {
            return Err(Error::ChainId(format!(
                "Current chain ID: {}, Tendermint chain ID: {}",
                chain_id, init.chain_id
            )));
        }
        if crate::migrating_state().is_some() {
            let rsp = response::InitChain {
                validators: self
                    .get_abci_validator_updates(true, |pk, power| {
                        let pub_key: crate::tendermint::PublicKey = pk.into();
                        let power =
                            crate::tendermint::vote::Power::try_from(power)
                                .unwrap();
                        validator::Update { pub_key, power }
                    })
                    .expect("Must be able to set genesis validator set"),
                app_hash: self
                    .state
                    .in_mem()
                    .merkle_root()
                    .0
                    .to_vec()
                    .try_into()
                    .expect("Infallible"),
                ..Default::default()
            };
            debug_assert!(!rsp.validators.is_empty());
            debug_assert!(
                !Vec::<u8>::from(rsp.app_hash.clone())
                    .iter()
                    .all(|&b| b == 0)
            );
            return Ok(rsp);
        }

        // Read the genesis files
        #[cfg(not(any(test, fuzzing, feature = "benches")))]
        let genesis = {
            let chain_dir = self.base_dir.join(chain_id);
            genesis::chain::Finalized::read_toml_files(&chain_dir)
                .expect("Missing or invalid genesis files")
        };
        #[cfg(any(test, fuzzing, feature = "benches"))]
        let genesis = {
            let chain_dir = self.base_dir.join(chain_id);
            if chain_dir.join(genesis::chain::METADATA_FILE_NAME).exists() {
                genesis::chain::Finalized::read_toml_files(&chain_dir)
                    .expect("Missing or invalid genesis files")
            } else {
                genesis::make_dev_genesis(num_validators, &chain_dir)
            }
        };

        let mut validation = InitChainValidation::new(self, false);
        let _ = validation.run(
            init,
            genesis,
            #[cfg(any(test, feature = "testing"))]
            num_validators,
        );
        // propagate errors or panic
        validation.error_out()?;

        // Init masp commitment tree and anchor
        let empty_commitment_tree: CommitmentTree<Node> =
            CommitmentTree::empty();
        let anchor = empty_commitment_tree.root();
        let note_commitment_tree_key =
            token::storage_key::masp_commitment_tree_key();
        self.state
            .write(&note_commitment_tree_key, empty_commitment_tree)
            .unwrap();
        let commitment_tree_anchor_key =
            token::storage_key::masp_commitment_anchor_key(anchor);
        self.state.write(&commitment_tree_anchor_key, ()).unwrap();

        // Init masp convert anchor
        let convert_anchor_key = token::storage_key::masp_convert_anchor_key();
        self.state.write(
            &convert_anchor_key,
            namada_sdk::hash::Hash(
                bls12_381::Scalar::from(
                    self.state.in_mem().conversion_state.tree.root(),
                )
                .to_bytes_le(),
            ),
        )?;

        // Set the initial validator set
        response.validators = self
            .get_abci_validator_updates(true, |pk, power| {
                let pub_key: crate::tendermint::PublicKey = pk.into();
                let power =
                    crate::tendermint::vote::Power::try_from(power).unwrap();
                validator::Update { pub_key, power }
            })
            .expect("Must be able to set genesis validator set");
        debug_assert!(!response.validators.is_empty());
        Ok(response)
    }
}

impl<D, H> InitChainValidation<'_, D, H>
where
    D: DB + for<'iter> DBIter<'iter> + Sync + 'static,
    H: StorageHasher + Sync + 'static,
{
    pub fn run(
        &mut self,
        init: request::InitChain,
        genesis: genesis::chain::Finalized,
        #[cfg(any(test, feature = "testing"))] _num_validators: u64,
    ) -> ControlFlow<()> {
        let ts: protobuf::Timestamp = init.time.into();
        let initial_height = init.initial_height.into();
        // TODO(informalsystems/tendermint-rs#870): hacky conversion
        let genesis_time: DateTimeUtc = (Utc.timestamp_opt(
            ts.seconds,
            u32::try_from(ts.nanos).expect("Time nanos cannot be negative"),
        ))
        .single()
        .expect("genesis time should be a valid timestamp")
        .into();

        // Initialize protocol parameters
        let parameters = genesis.get_chain_parameters(&self.wasm_dir);
        self.store_wasms(&parameters)?;
        parameters::init_storage(&parameters, &mut self.state).unwrap();

        // Initialize governance parameters
        let gov_params = genesis.get_gov_params();
        gov_params.init_storage(&mut self.state).unwrap();

        // configure the Ethereum bridge if the configuration is set.
        if let Some(config) = genesis.get_eth_bridge_params() {
            tracing::debug!("Initializing Ethereum bridge storage.");
            config.init_storage(&mut self.state);
            self.update_eth_oracle(&Default::default());
        } else {
            self.state
                .write(
                    &eth_bridge::storage::active_key(),
                    EthBridgeStatus::Disabled,
                )
                .unwrap();
        }

        // Initialize IBC parameters
        let ibc_params = genesis.get_ibc_params();
        ibc_params.init_storage(&mut self.state).unwrap();

        // Depends on parameters being initialized
        self.state
            .in_mem_mut()
            .init_genesis_epoch(initial_height, genesis_time, &parameters)
            .expect("Initializing genesis epoch must not fail");

        // PoS system depends on epoch being initialized
        let pos_params = genesis.get_pos_params();
        let (current_epoch, _gas) = self.state.in_mem().get_current_epoch();
        proof_of_stake::init_genesis(
            &mut self.state,
            &pos_params,
            current_epoch,
        )
        .expect("Must be able to initialize PoS genesis storage");

        // PGF parameters
        let pgf_params = genesis.get_pgf_params();
        pgf_params
            .init_storage(&mut self.state)
            .expect("Should be able to initialized PGF at genesis");

        // Loaded VP code cache to avoid loading the same files multiple times
        let mut vp_cache: HashMap<String, Vec<u8>> = HashMap::default();
        self.init_token_accounts(&genesis);
        let _ = self.init_token_balances(&genesis);
        let _ =
            self.apply_genesis_txs_established_account(&genesis, &mut vp_cache);
        let _ = self.apply_genesis_txs_validator_account(
            &genesis,
            &mut vp_cache,
            &pos_params,
            current_epoch,
        );
        self.apply_genesis_txs_bonds(&genesis);

        proof_of_stake::compute_and_store_total_consensus_stake::<
            _,
            governance::Store<_>,
        >(&mut self.state, current_epoch)
        .expect("Could not compute total consensus stake at genesis");
        // This has to be done after `apply_genesis_txs_validator_account`
        proof_of_stake::copy_genesis_validator_sets::<_, governance::Store<_>>(
            &mut self.state,
            &pos_params,
            current_epoch,
        )
        .expect("Must be able to copy PoS genesis validator sets");

        ibc::init_genesis_storage(&mut self.state);
        ControlFlow::Continue(())
    }

    /// Look-up WASM code of a genesis VP by its name
    fn lookup_vp(
        &mut self,
        name: &str,
        genesis: &genesis::chain::Finalized,
        vp_cache: &mut HashMap<String, Vec<u8>>,
    ) -> ControlFlow<(), Vec<u8>> {
        use namada_sdk::collections::hash_map::Entry;
        let Some(vp_filename) = self
            .validate(
                genesis
                    .vps
                    .wasm
                    .get(name)
                    .map(|conf| conf.filename.clone())
                    .ok_or_else(|| {
                        Panic::MissingVpWasmConfig(name.to_string())
                    }),
            )
            .or_placeholder(None)?
        else {
            return self.proceed_with(vec![]);
        };
        let code = match vp_cache.entry(vp_filename.clone()) {
            Entry::Occupied(o) => o.get().clone(),
            Entry::Vacant(v) => {
                let code = self
                    .validate(
                        wasm_loader::read_wasm(&self.wasm_dir, &vp_filename)
                            .map_err(|e| {
                                Panic::ReadingWasm(vp_filename, e.to_string())
                            }),
                    )
                    .or_placeholder(Some(vec![]))?
                    .unwrap();
                v.insert(code).clone()
            }
        };
        self.proceed_with(code)
    }

    fn store_wasms(&mut self, params: &Parameters) -> ControlFlow<()> {
        let Parameters {
            tx_allowlist,
            vp_allowlist,
            implicit_vp_code_hash,
            ..
        } = params;
        let mut is_implicit_vp_stored = false;

        let Some(checksums) = self
            .validate(
                wasm_loader::Checksums::read_checksums(&self.wasm_dir)
                    .map_err(|_| Panic::ChecksumsFile),
            )
            .or_placeholder(None)?
        else {
            return self.proceed_with(());
        };

        for (name, full_name) in checksums.0.iter() {
            let code = self
                .validate(
                    wasm_loader::read_wasm(&self.wasm_dir, name)
                        .map_err(Error::ReadingWasm),
                )
                .or_placeholder(Some(vec![]))?
                .unwrap();

            let code_hash = CodeHash::sha256(&code);
            let code_len = self
                .validate(
                    u64::try_from(code.len())
                        .map_err(|e| Error::LoadingWasm(e.to_string())),
                )
                .or_placeholder(Some(1))?
                .unwrap();

            let elements = full_name.split('.').collect::<Vec<&str>>();
            let checksum = self
                .validate(
                    elements
                        .get(1)
                        .map(|c| c.to_string().to_uppercase())
                        .ok_or_else(|| {
                            Error::LoadingWasm(format!(
                                "invalid full name: {}",
                                full_name
                            ))
                        }),
                )
                .or_placeholder(Some(code_hash.to_string()))?
                .unwrap();

            self.validate(if checksum == code_hash.to_string() {
                Ok(())
            } else {
                Err(Panic::Checksum(name.to_string()))
            })
            .or_placeholder(None)?;

            if (tx_allowlist.is_empty() && vp_allowlist.is_empty())
                || tx_allowlist.contains(&code_hash.to_string().to_lowercase())
                || vp_allowlist.contains(&code_hash.to_string().to_lowercase())
            {
                self.validate(
                    validate_untrusted_wasm(&code)
                        .map_err(|e| Error::LoadingWasm(e.to_string())),
                )
                .or_placeholder(None)?;

                #[cfg(not(any(test, fuzzing)))]
                if name.starts_with("tx_") {
                    self.tx_wasm_cache.pre_compile(
                        &code,
                        namada_sdk::gas::GasMeterKind::MutGlobal,
                    );
                } else if name.starts_with("vp_") {
                    self.vp_wasm_cache.pre_compile(
                        &code,
                        namada_sdk::gas::GasMeterKind::MutGlobal,
                    );
                }

                let code_key = Key::wasm_code(&code_hash);
                let code_len_key = Key::wasm_code_len(&code_hash);
                let hash_key = Key::wasm_hash(name);
                let code_name_key = Key::wasm_code_name(name.to_owned());

                self.state.write(&code_key, code).unwrap();
                self.state.write(&code_len_key, code_len).unwrap();
                self.state.write(&hash_key, code_hash).unwrap();
                if &Some(code_hash) == implicit_vp_code_hash {
                    is_implicit_vp_stored = true;
                }
                self.state.write(&code_name_key, code_hash).unwrap();
            } else {
                tracing::warn!("The wasm {name} isn't allowed.");
                self.warn(Warning::DisallowedWasm(name.to_string()));
            }
        }

        // check if implicit_vp wasm is stored
        if !is_implicit_vp_stored {
            self.register_err(Panic::MissingImplicitVP(
                match implicit_vp_code_hash {
                    None => "None".to_string(),
                    Some(h) => h.to_string(),
                },
            ));
        }

        self.proceed_with(())
    }

    /// Init genesis token accounts
    fn init_token_accounts(&mut self, genesis: &genesis::chain::Finalized) {
        let mut token_map = BTreeMap::new();
        let native_alias = &genesis.parameters.parameters.native_token;
        for (alias, token) in &genesis.tokens.token {
            tracing::debug!("Initializing token {alias}");

            let FinalizedTokenConfig {
                address,
                config: TokenConfig { denom, masp_params },
            } = token;
            if alias == native_alias && masp_params.is_none() {
                self.register_err(Panic::MissingMaspParams);
            }
            // associate a token with its denomination.
            write_denom(&mut self.state, address, *denom).unwrap();
            namada_sdk::token::write_params(
                masp_params,
                &mut self.state,
                address,
                denom,
            )
            .unwrap();
            if masp_params.is_some() {
                // add token addresses to the masp reward conversions lookup
                // table.
                let alias = alias.to_string();
                token_map.insert(alias, address.clone());
            }
        }
        self.state
            .write(&masp_token_map_key(), token_map)
            .expect("Couldn't init token accounts");
    }

    /// Init genesis token balances
    fn init_token_balances(
        &mut self,
        genesis: &genesis::chain::Finalized,
    ) -> ControlFlow<()> {
        for (token_alias, TokenBalances(balances)) in &genesis.balances.token {
            tracing::debug!("Initializing token balances {token_alias}");

            let Some(token_address) = self
                .validate(
                    genesis
                        .tokens
                        .token
                        .get(token_alias)
                        .ok_or_else(|| {
                            Panic::MissingTokenConfig(token_alias.to_string())
                        })
                        .map(|conf| &conf.address),
                )
                .or_placeholder(None)?
            else {
                continue;
            };

            for (owner, balance) in balances {
                tracing::info!(
                    "Crediting {} {} tokens to {}",
                    balance,
                    token_alias,
                    owner,
                );
                credit_tokens(
                    &mut self.state,
                    token_address,
                    owner,
                    balance.amount(),
                )
                .expect("Couldn't credit initial balance");
            }
        }
        self.proceed_with(())
    }

    /// Apply genesis txs to initialize established accounts
    fn apply_genesis_txs_established_account(
        &mut self,
        genesis: &genesis::chain::Finalized,
        vp_cache: &mut HashMap<String, Vec<u8>>,
    ) -> ControlFlow<()> {
        if let Some(txs) = genesis.transactions.established_account.as_ref() {
            for FinalizedEstablishedAccountTx {
                address,
                tx:
                    EstablishedAccountTx {
                        vp,
                        threshold,
                        public_keys,
                    },
            } in txs
            {
                tracing::debug!(
                    "Applying genesis tx to init an established account \
                     {address}"
                );
                let vp_code = self.lookup_vp(vp, genesis, vp_cache)?;
                let code_hash = CodeHash::sha256(&vp_code);
                self.state
                    .write(&Key::validity_predicate(address), code_hash)
                    .unwrap();

                let public_keys: Vec<_> =
                    public_keys.iter().map(|pk| pk.raw.clone()).collect();

                namada_sdk::account::init_account_storage(
                    &mut self.state,
                    address,
                    &public_keys,
                    *threshold,
                )
                .unwrap();

                for pk in &public_keys {
                    let implicit_addr = pk.into();

                    namada_sdk::account::init_account_storage(
                        &mut self.state,
                        &implicit_addr,
                        std::slice::from_ref(pk),
                        1,
                    )
                    .unwrap();
                }
            }
        }
        self.proceed_with(())
    }

    /// Apply genesis txs to initialize validator accounts
    fn apply_genesis_txs_validator_account(
        &mut self,
        genesis: &genesis::chain::Finalized,
        vp_cache: &mut HashMap<String, Vec<u8>>,
        params: &PosParams,
        current_epoch: namada_sdk::chain::Epoch,
    ) -> ControlFlow<()> {
        if let Some(txs) = genesis.transactions.validator_account.as_ref() {
            for FinalizedValidatorAccountTx {
                tx:
                    SignedTx {
                        data:
                            ValidatorAccountTx {
                                address,
                                vp,
                                commission_rate,
                                max_commission_rate_change,
                                metadata,
                                net_address: _,
                                consensus_key,
                                protocol_key,
                                tendermint_node_key: _,
                                eth_hot_key,
                                eth_cold_key,
                                ..
                            },
                        ..
                    },
            } in txs
            {
                let address = &Address::Established(address.raw.clone());

                tracing::debug!(
                    "Applying genesis tx to init a validator account {address}"
                );

                let vp_code = self.lookup_vp(vp, genesis, vp_cache)?;
                let code_hash = CodeHash::sha256(&vp_code);
                self.state
                    .write(&Key::validity_predicate(address), code_hash)
                    .expect("Unable to write user VP");

                self.state
                    .write(&protocol_pk_key(address), &protocol_key.pk.raw)
                    .expect("Unable to set genesis user protocol public key");

                if let Err(err) =
                    proof_of_stake::become_validator::<_, governance::Store<_>>(
                        &mut self.state,
                        BecomeValidator {
                            params,
                            address,
                            consensus_key: &consensus_key.pk.raw,
                            protocol_key: &protocol_key.pk.raw,
                            eth_cold_key: &eth_cold_key.pk.raw,
                            eth_hot_key: &eth_hot_key.pk.raw,
                            current_epoch,
                            commission_rate: *commission_rate,
                            max_commission_rate_change:
                                *max_commission_rate_change,
                            metadata: metadata.clone(),
                            offset_opt: Some(0),
                        },
                    )
                {
                    tracing::warn!(
                        "Genesis init genesis validator tx for {address} \
                         failed with {err}. Skipping."
                    );
                    self.warn(Warning::Validator(
                        address.to_string(),
                        err.to_string(),
                    ));
                    continue;
                }
            }
        }
        self.proceed_with(())
    }

    /// Apply genesis txs to transfer tokens
    fn apply_genesis_txs_bonds(&mut self, genesis: &genesis::chain::Finalized) {
        let (current_epoch, _gas) = self.state.in_mem().get_current_epoch();
        if let Some(txs) = &genesis.transactions.bond {
            for BondTx {
                source,
                validator,
                amount,
                ..
            } in txs
            {
                tracing::debug!(
                    "Applying genesis tx to bond {} native tokens from \
                     {source} to {validator}",
                    amount,
                );

                if let Err(err) = proof_of_stake::bond_tokens::<
                    _,
                    governance::Store<_>,
                    token::Store<_>,
                >(
                    &mut self.state,
                    Some(&source.address()),
                    validator,
                    amount.amount(),
                    current_epoch,
                    Some(0),
                ) {
                    tracing::warn!(
                        "Genesis bond tx failed with: {err}. Skipping."
                    );
                    self.warn(Warning::FailedBond(
                        source.to_string(),
                        validator.to_string(),
                        *amount,
                        err.to_string(),
                    ));
                    continue;
                };
            }
        }
    }
}

/// A helper struct to accumulate errors in genesis files while
/// attempting to initialize the ledger
#[derive(Debug)]
pub struct InitChainValidation<'shell, D, H>
where
    D: DB + for<'iter> DBIter<'iter> + Sync + 'static,
    H: StorageHasher + Sync + 'static,
{
    /// Errors that can be encountered while initializing chain
    /// and are propagated up the stack in normal flow. Ultimately
    /// these are reported back to Comet BFT
    errors: Vec<Error>,
    /// Errors that cause `init_chain` to panic in normal flow but are not
    /// `expect` calls, so they could reasonably occur. These are demoted
    /// to errors while validating correctness of genesis files pre-network
    /// launch.
    panics: Vec<Panic>,
    /// Events that should not occur but would not prevent the chain from
    /// being successfully initialized. However, we don't reasonably expect
    /// to get any as these are checked as part of validating genesis
    /// templates.
    warnings: Vec<Warning>,
    dry_run: bool,
    shell: &'shell mut Shell<D, H>,
}

impl<D, H> std::ops::Deref for InitChainValidation<'_, D, H>
where
    D: DB + for<'iter> DBIter<'iter> + Sync + 'static,
    H: StorageHasher + Sync + 'static,
{
    type Target = Shell<D, H>;

    fn deref(&self) -> &Self::Target {
        self.shell
    }
}

impl<D, H> std::ops::DerefMut for InitChainValidation<'_, D, H>
where
    D: DB + for<'iter> DBIter<'iter> + Sync + 'static,
    H: StorageHasher + Sync + 'static,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.shell
    }
}

impl<'shell, D, H> InitChainValidation<'shell, D, H>
where
    D: DB + for<'iter> DBIter<'iter> + Sync + 'static,
    H: StorageHasher + Sync + 'static,
{
    pub fn new(
        shell: &'shell mut Shell<D, H>,
        dry_run: bool,
    ) -> InitChainValidation<'shell, D, H> {
        Self {
            shell,
            errors: vec![],
            panics: vec![],
            warnings: vec![],
            dry_run,
        }
    }

    pub fn run_validation(
        &mut self,
        chain_id: String,
        genesis: config::genesis::chain::Finalized,
    ) {
        use crate::tendermint::block::Size;
        use crate::tendermint::consensus::Params;
        use crate::tendermint::consensus::params::ValidatorParams;
        use crate::tendermint::evidence::{Duration, Params as Evidence};
        use crate::tendermint::time::Time;

        // craft a request to initialize the chain
        let init = request::InitChain {
            time: Time::now(),
            chain_id,
            consensus_params: Params {
                block: Size {
                    max_bytes: 0,
                    max_gas: 0,
                    time_iota_ms: 0,
                },
                evidence: Evidence {
                    max_age_num_blocks: 0,
                    max_age_duration: Duration(Default::default()),
                    max_bytes: 0,
                },
                validator: ValidatorParams {
                    pub_key_types: vec![],
                },
                version: None,
                abci: Default::default(),
            },
            validators: vec![],
            app_state_bytes: Default::default(),
            initial_height: 0u32.into(),
        };
        let _ = self.run(
            init,
            genesis,
            #[cfg(any(test, feature = "testing"))]
            1,
        );
    }

    /// Print out a report of errors encountered while dry-running
    /// genesis files
    pub fn report(&self) {
        use color_eyre::owo_colors::{OwoColorize, Style};
        let separator: String = ["="; 60].into_iter().collect();
        println!(
            "\n\n{}\n{}\n{}\n\n",
            separator,
            "Report".bold().underline(),
            separator
        );
        if self.errors.is_empty()
            && self.panics.is_empty()
            && self.warnings.is_empty()
        {
            println!(
                "{}\n",
                "Genesis files were dry-run successfully"
                    .bright_green()
                    .underline()
            );
            return;
        }

        if !self.warnings.is_empty() {
            println!("{}\n\n", "Warnings".yellow().underline());
            let warnings = Style::new().yellow();
            for warning in &self.warnings {
                println!("{}\n", warning.to_string().style(warnings));
            }
        }

        if !self.errors.is_empty() {
            println!("{}\n\n", "Errors".magenta().underline());
            let errors = Style::new().magenta();
            for error in &self.errors {
                println!("{}\n", error.to_string().style(errors));
            }
        }

        if !self.panics.is_empty() {
            println!("{}\n\n", "Panics".bright_red().underline());
            let panics = Style::new().bright_red();
            for panic in &self.panics {
                println!("{}\n", panic.to_string().style(panics));
            }
        }
    }

    /// Add a warning
    fn warn(&mut self, warning: Warning) {
        self.warnings.push(warning);
    }

    /// Categorize an error as normal or something that would panic.
    fn register_err<E: Into<ErrorType>>(&mut self, err: E) {
        match err.into() {
            ErrorType::Runtime(e) => self.errors.push(e),
            ErrorType::DryRun(e) => self.panics.push(e),
        }
    }

    /// Categorize the error encountered and return a handle to allow
    /// the code to specify how to proceed.
    fn validate<T, E>(&mut self, res: std::result::Result<T, E>) -> Policy<T>
    where
        E: Into<ErrorType>,
    {
        match res {
            Ok(data) => Policy {
                result: Some(data),
                dry_run: self.dry_run,
            },
            Err(e) => {
                self.register_err(e);
                Policy {
                    result: None,
                    dry_run: self.dry_run,
                }
            }
        }
    }

    /// Check if any errors have been encountered
    fn is_ok(&self) -> bool {
        self.errors.is_empty() && self.panics.is_empty()
    }

    /// This should only be called after checking that `is_ok` returned false.
    fn error_out(mut self) -> ShellResult<()> {
        if self.is_ok() {
            return Ok(());
        }
        if !self.panics.is_empty() {
            panic!(
                "Namada ledger failed to initialize due to: {}",
                self.panics.remove(0)
            );
        } else {
            Err(self.errors.remove(0))
        }
    }

    /// Used to indicate to the functions up the stack to begin panicking
    /// if not dry running a genesis file
    fn proceed_with<T>(&self, value: T) -> ControlFlow<(), T> {
        if self.dry_run || self.is_ok() {
            ControlFlow::Continue(value)
        } else {
            ControlFlow::Break(())
        }
    }
}

enum ErrorType {
    Runtime(Error),
    DryRun(Panic),
}

impl From<Error> for ErrorType {
    fn from(err: Error) -> Self {
        Self::Runtime(err)
    }
}

impl From<Panic> for ErrorType {
    fn from(err: Panic) -> Self {
        Self::DryRun(err)
    }
}

/// Changes the control flow of `init_chain` depending on whether
/// or not it is a dry-run. If so, errors / panics are accumulated to make
/// a report rather than immediately exiting.
struct Policy<T> {
    result: Option<T>,
    dry_run: bool,
}

impl<T> Policy<T> {
    /// A default value to return if an error / panic is encountered
    /// during a dry-run. This allows `init_chain` to continue.
    fn or_placeholder(self, value: Option<T>) -> ControlFlow<(), Option<T>> {
        if let Some(data) = self.result {
            ControlFlow::Continue(Some(data))
        } else if self.dry_run {
            ControlFlow::Continue(value)
        } else {
            ControlFlow::Break(())
        }
    }
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use namada_apps_lib::wallet::defaults;
    use namada_sdk::string_encoding::StringEncoded;
    use namada_sdk::wallet::alias::Alias;

    use super::*;
    use crate::config::genesis::{GenesisAddress, transactions};
    use crate::shell::test_utils::TestShell;

    /// Test that the init-chain handler never commits changes directly to the
    /// DB.
    #[test]
    fn test_init_chain_doesnt_commit_db() {
        let (shell, _recv, _, _) = test_utils::setup();

        // Collect all storage key-vals into a sorted map
        let store_block_state = |shell: &TestShell| -> BTreeMap<_, _> {
            shell
                .state
                .db()
                .iter_prefix(None)
                .map(|(key, val, _gas)| (key, val))
                .collect()
        };

        // Store the full state in sorted map
        let initial_storage_state: std::collections::BTreeMap<String, Vec<u8>> =
            store_block_state(&shell);

        // Store the full state again
        let storage_state: std::collections::BTreeMap<String, Vec<u8>> =
            store_block_state(&shell);

        // The storage state must be unchanged
        itertools::assert_equal(
            initial_storage_state.iter(),
            storage_state.iter(),
        );
    }

    /// Tests validation works properly on `lookup_vp`.
    /// This function can fail if
    /// *the wasm requested has no config in the genesis files
    /// * cannot be read from disk.
    #[test]
    fn test_dry_run_lookup_vp() {
        let (mut shell, _x, _y, _z) = TestShell::new_at_height(0);
        shell.wasm_dir = PathBuf::new();
        let mut genesis = genesis::make_dev_genesis(1, &shell.base_dir);
        let mut initializer = InitChainValidation::new(&mut shell, true);

        let mut vp_cache = HashMap::new();
        let code = initializer.lookup_vp("vp_user", &genesis, &mut vp_cache);
        assert_eq!(code, ControlFlow::Continue(vec![]));
        assert_eq!(
            *vp_cache.get("vp_user.wasm").expect("Test failed"),
            Vec::<u8>::new()
        );
        let [Panic::ReadingWasm(_, _)]: [Panic; 1] =
            initializer.panics.clone().try_into().expect("Test failed")
        else {
            panic!("Test failed")
        };

        initializer.panics.clear();
        genesis.vps.wasm.remove("vp_user").expect("Test failed");
        let code = initializer.lookup_vp("vp_user", &genesis, &mut vp_cache);
        assert_eq!(code, ControlFlow::Continue(vec![]));
        let [Panic::MissingVpWasmConfig(_)]: [Panic; 1] =
            initializer.panics.clone().try_into().expect("Test failed")
        else {
            panic!("Test failed")
        };
    }

    /// Test validation of `store_wasms`.
    /// This can fail if
    /// * The checksums file cannot be found.
    /// * A wasm file in the checksums file cannot be read from disk
    /// * A checksum entry is invalid
    /// * A wasm's code hash does not match it's checksum entry
    /// * the wasm isn't allowed
    /// * no vp_implicit wasm is stored
    #[test]
    fn test_dry_run_store_wasms() {
        let (mut shell, _x, _y, _z) = TestShell::new_at_height(0);
        let test_dir = tempfile::tempdir().unwrap();
        shell.wasm_dir = test_dir.path().into();

        let genesis = genesis::make_dev_genesis(1, &shell.base_dir);
        let mut initializer = InitChainValidation::new(&mut shell, true);

        let res = initializer
            .store_wasms(&genesis.get_chain_parameters(PathBuf::new()));
        assert_eq!(res, ControlFlow::Continue(()));
        let expected = vec![Panic::ChecksumsFile];
        assert_eq!(expected, initializer.panics);
        initializer.panics.clear();

        let checksums_file = test_dir.path().join("checksums.json");
        std::fs::write(
            &checksums_file,
            r#"{
            "tx_get_rich.wasm": "tx_get_rich.moneymoneymoney"
        }"#,
        )
        .expect("Test failed");
        let res = initializer
            .store_wasms(&genesis.get_chain_parameters(test_dir.path()));
        assert_eq!(res, ControlFlow::Continue(()));
        let errors = initializer.errors.iter().collect::<Vec<_>>();
        let [Error::ReadingWasm(_), Error::LoadingWasm(_)]: [&Error; 2] =
            errors.try_into().expect("Test failed")
        else {
            panic!("Test failed");
        };
        let expected_panics = vec![
            Panic::Checksum("tx_get_rich.wasm".into()),
            Panic::MissingImplicitVP("None".into()),
        ];
        assert_eq!(initializer.panics, expected_panics);

        initializer.panics.clear();
        initializer.errors.clear();

        std::fs::write(
            checksums_file,
            r#"{
            "tx_stuff.wasm": "tx_stuff"
        }"#,
        )
        .expect("Test failed");
        let res = initializer
            .store_wasms(&genesis.get_chain_parameters(test_dir.path()));
        assert_eq!(res, ControlFlow::Continue(()));
        let errors = initializer.errors.iter().collect::<Vec<_>>();
        let [
            Error::ReadingWasm(_),
            Error::LoadingWasm(_),
            Error::LoadingWasm(_),
        ]: [&Error; 3] = errors.try_into().expect("Test failed")
        else {
            panic!("Test failed");
        };
        let expected_panics = vec![Panic::MissingImplicitVP("None".into())];
        assert_eq!(initializer.panics, expected_panics);
    }

    /// Test validation of `init_token_balance`.
    /// This can fail if a token alias with no
    /// corresponding config is encountered.
    #[test]
    fn test_dry_run_init_token_balance() {
        let (mut shell, _x, _y, _z) = TestShell::new_at_height(0);
        shell.wasm_dir = PathBuf::new();
        let mut genesis = genesis::make_dev_genesis(1, &shell.base_dir);
        let mut initializer = InitChainValidation::new(&mut shell, true);
        let token_alias = Alias::from_str("apfel").unwrap();
        genesis
            .tokens
            .token
            .remove(&token_alias)
            .expect("Test failed");
        let res = initializer.init_token_balances(&genesis);
        assert_eq!(res, ControlFlow::Continue(()));
        let [Panic::MissingTokenConfig(_)]: [Panic; 1] =
            initializer.panics.clone().try_into().expect("Test failed")
        else {
            panic!("Test failed")
        };
    }

    /// Test validation of `apply_genesis_txs_bonds`
    /// This can fail for
    ///  * insufficient funds
    /// * bonding to a non-validator
    #[test]
    fn test_dry_run_genesis_bonds() {
        let (mut shell, _x, _y, _z) = TestShell::new_at_height(0);
        shell.wasm_dir = PathBuf::new();
        let mut genesis = genesis::make_dev_genesis(1, &shell.base_dir);
        let mut initializer = InitChainValidation::new(&mut shell, true);
        let default_addresses: HashMap<Alias, Address> =
            defaults::addresses().into_iter().collect();
        let albert_address = if let Some(Address::Established(albert)) =
            default_addresses.get(&Alias::from_str("albert").unwrap())
        {
            albert.clone()
        } else {
            panic!("Test failed")
        };
        // Initialize governance parameters
        let gov_params = genesis.get_gov_params();
        gov_params.init_storage(&mut initializer.state).unwrap();
        // PoS system depends on epoch being initialized
        let pos_params = genesis.get_pos_params();
        let (current_epoch, _gas) =
            initializer.state.in_mem().get_current_epoch();
        proof_of_stake::init_genesis(
            &mut initializer.state,
            &pos_params,
            current_epoch,
        )
        .expect("Must be able to initialize PoS genesis storage");

        genesis.transactions.bond = Some(vec![transactions::BondTx {
            source: GenesisAddress::EstablishedAddress(albert_address.clone()),
            validator: defaults::albert_address(),
            amount: token::DenominatedAmount::new(
                token::Amount::from_uint(1, 6).unwrap(),
                6.into(),
            ),
        }]);

        // bonds should fail since no balances have been initialized
        let albert_address_str = StringEncoded::new(albert_address).to_string();
        initializer.apply_genesis_txs_bonds(&genesis);
        let expected = vec![Warning::FailedBond(
            albert_address_str.clone(),
            albert_address_str.clone(),
            token::DenominatedAmount::new(
                token::Amount::from_uint(1, 6).unwrap(),
                6.into(),
            ),
            format!("{albert_address_str} has insufficient balance"),
        )];
        assert_eq!(expected, initializer.warnings);
        initializer.warnings.clear();

        // initialize balances
        let res = initializer.init_token_balances(&genesis);
        assert_eq!(res, ControlFlow::Continue(()));

        initializer.apply_genesis_txs_bonds(&genesis);
        let expected = vec![Warning::FailedBond(
            albert_address_str.clone(),
            albert_address_str.clone(),
            token::DenominatedAmount::new(
                token::Amount::from_uint(1, 6).unwrap(),
                6.into(),
            ),
            format!(
                "The given address {} is not a validator address",
                albert_address_str
            ),
        )];
        assert_eq!(expected, initializer.warnings);
    }

    #[test]
    fn test_dry_run_native_token_masp_params() {
        let (mut shell, _x, _y, _z) = TestShell::new_at_height(0);
        shell.wasm_dir = PathBuf::new();
        let mut genesis = genesis::make_dev_genesis(1, &shell.base_dir);
        let mut initializer = InitChainValidation::new(&mut shell, true);
        genesis
            .tokens
            .token
            .get_mut(&genesis.parameters.parameters.native_token)
            .expect("Test failed")
            .config
            .masp_params = None;
        initializer.init_token_accounts(&genesis);
        let [panic]: [Panic; 1] =
            initializer.panics.clone().try_into().expect("Test failed");
        assert_eq!(panic, Panic::MissingMaspParams);
    }
}
